# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: leon

import os

from absl import app, flags

import pandas
import matplotlib
matplotlib.use('Agg')
import tabulate

FLAGS = flags.FLAGS


def calc_edge(bid_x, ask_x, bid_y, ask_y, basis):
  gap = bid_x - ask_y + basis
  mid = (bid_x + ask_x + bid_y + ask_y) / 4
  edge_bp = gap / mid * 10000
  return edge_bp


def calc_arb_stats(df, exchange, product):
  headers = [
      'exchange_pair',
      'product',
      'basis_ma',
      'edge_bp',
      'fluctuate_count',
      'Kraken->%s_count' % exchange,
      '%s->Kraken_count' % exchange
  ]
  for edge_bp in [0, 2, 5, 10, 15, 20]:
    for basis_ma in [5, 15, 30]:
      fluctuate_count = 0
      x_to_y_count = 0
      y_to_x_count = 0
      continuous_x_to_y_count = 0  # 10 times for a cycle
      continuous_y_to_x_count = 0  # 10 times for a cycle
      for _, row in df.iterrows():
        curr_edge_x_to_y = calc_edge(row['true_bid_x'],
                                     row['true_ask_x'],
                                     row['true_bid_y'],
                                     row['true_ask_y'],
                                     -row['mid_basis_%d' % basis_ma])
        curr_edge_y_to_x = calc_edge(row['true_bid_y'],
                                     row['true_ask_y'],
                                     row['true_bid_x'],
                                     row['true_ask_x'],
                                     row['mid_basis_%d' % basis_ma])
        if curr_edge_x_to_y > edge_bp:
          x_to_y_count += 1
          continuous_x_to_y_count += 1
          continuous_y_to_x_count = 0
        elif curr_edge_y_to_x > edge_bp:
          y_to_x_count += 1
          continuous_y_to_x_count += 1
          continuous_x_to_y_count = 0

        if continuous_x_to_y_count == 10:
          continuous_x_to_y_count += 1
          fluctuate_count += 1
        elif continuous_y_to_x_count == 10:
          continuous_y_to_x_count += 1
          fluctuate_count += 1

      row = [
          'Kraken/%s' % exchange,
          product,
          basis_ma,
          edge_bp,
          fluctuate_count,
          x_to_y_count,
          y_to_x_count
      ]
      print(tabulate.tabulate([row], headers=headers))
      print()


def calculate_basis(df):
  df['mid_diff'] = (df['true_ask_x'] + df['true_bid_x'] - df['true_ask_y'] - df['true_bid_y']) / 2
  for win_size in [5, 15, 30]:
    df['mid_basis_%d' % win_size] = \
        df['mid_diff'].rolling(window=win_size * 60, min_periods=1).mean()
  return df


def merge_dfs(df1, df2):
  return pandas.merge(df1, df2, how='inner', left_index=True, right_index=True)


def sample_df_per_second(df):
  return df.resample('S', closed='left', label='left').last()


def read_csv_into_df(csv_root, csv):
  csv_dir = os.path.join(csv_root, csv)
  df = pandas.read_csv(csv_dir, sep=',', header=0)
  df['datetime'] = pandas.to_datetime(df['timestamp'], unit='ns')
  df = df.set_index('datetime')
  return df


def main(argv):
  csv_root = FLAGS.csv_root
  assert csv_root, '--csv_root must be specified.'

  products = ['BCH-USD', 'BTC-USD', 'EOS-USD', 'XRP-USD', 'ETH-USD', 'ETC-USD']
  for product in products:
    kraken_price_df = read_csv_into_df(csv_root, '%s.%s.csv' % (product, 'Kraken'))
    for exchange in ['Binance', 'Bitfinex', 'Huobi', 'Okex']:
      other_price_df = read_csv_into_df(csv_root, '%sT.%s.csv' % (product, exchange))
      sampled_df1 = sample_df_per_second(kraken_price_df)
      sampled_df2 = sample_df_per_second(other_price_df)
      merged_df = merge_dfs(sampled_df1, sampled_df2)
      merged_df = calculate_basis(merged_df)
      calc_arb_stats(merged_df, exchange, product)

  return 0


if __name__ == '__main__':
  flags.DEFINE_string('csv_root', None, 'Input csv files root directory.')

  app.run(main)
